#! /usr/bin/env python
# coding=utf-8

from Bio import SeqIO

import argparse
import sys
from argparse import RawTextHelpFormatter




parser = argparse.ArgumentParser(
    description='''
    通过表达量(一个bowtie结果) 定义miRNA 哪个 side 是成熟体 
    用法:
    define_mature_side_using_reads.py -m result.fasta -b all_map -o side_prop
    ''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-m',
                help='必须给定，输入的miRNA的fasta文件')

parser.add_argument('-b',
                help='必须给定，输出的bowtie结果')

parser.add_argument('-o',
                help='必须给定，输出的fasta文件 ')


args = parser.parse_args()

if not args.m or not args.b or not args.o:
    parser.print_help()
    sys.exit()


# 通过 表达量 判断miRNA 哪一边是 mature
miRNA_file = args.m


# 总的表达文件
bowtie_file = args.b

# 这里有个 parent_lista 文件 用于制定 预先选定的 成熟体 判定

#parent_file = r'parent_lista'

outfile  = args.o




mature_dic ={}

class miRNA(object):
	def __init__(self, seq):
		self.name = str(seq.id)
		self.seq = str(seq.seq)
		self.length = len(seq.seq)
		self.cov_5p_total = 0
		self.cov_3p_total = 0
		self.side = '5p'
	def find_with_side_is_mature(self):
		if self.cov_3p_total > self.cov_5p_total:
			self.side = '3p'


prop_dic = {}

# with open(parent_file) as fila:
# 	for i in fila:
# 		k = i.strip().split('\t')
# 		if len(k)>1:
# 			prop_dic[k[0]] = k[1]



miRNA_dic =  {}

for i in SeqIO.parse(miRNA_file,'fasta'):
	one_miRNA = miRNA(i)
	miRNA_dic[one_miRNA.name] = one_miRNA



with open(bowtie_file) as fila:
	for i in fila:
		k = i.split('#')[0].strip().split('\t')
		if len(k)>5 and k[1]=='+':
			target_name = k[2]
			if target_name in miRNA_dic:
				start = int(k[3])
				nub = 1
				if len(k[0].upper().split('_X'))>1:
					nub = int(k[0].upper().split('_X')[1])

				if start > miRNA_dic[target_name].length/2:

					miRNA_dic[target_name].cov_3p_total += nub
				else:
					miRNA_dic[target_name].cov_5p_total += nub




for i in miRNA_dic:
	miRNA_dic[i].find_with_side_is_mature()
	prop_dic[i] = miRNA_dic[i].side

with open(outfile,'w') as fila:
	for i in prop_dic:

		fila.write(i+'\t'+prop_dic[i])
		if i in miRNA_dic:
			fila.write('\t'+str(miRNA_dic[i].cov_5p_total)+'\t'+str(miRNA_dic[i].cov_3p_total))
		fila.write('\n')
	

				



				

